import math
import time
import numpy as np
from threading import Lock, Thread, Event

from ezgl import Matrix4x4
from ezrobot.utils.general import accurate_sleep
from .bezier_path import PoseBlender


class OnlineSegment():

    def __init__(self, blender: PoseBlender):
        self.num_poses = 0  # 控制点的数量 1 / 2 / 3
        self.blender = blender

        self.length_1 = 0   # 第一段的长度
        self.length_2 = 0   # 第二段的长度
        self.t_last = 0  # 已经走过的 t

        self.poses = [None, None, None]

        self.mutex = Lock()

    def addPose(self, pose):
        self.mutex.acquire()

        if self.num_poses == 0:
            self.poses[0] = pose

        elif self.num_poses == 1:
            self.poses[1] = pose
            self.length_1 = self.blender.distance(self.poses[0], self.poses[1])

        elif self.num_poses == 2:
            self.poses[0] = self.blender.interp(self.poses[0], self.poses[1], self.t_last)
            self.poses[2] = pose
            self.length_1 = self.blender.distance(self.poses[0], self.poses[1])
            self.length_2 = self.blender.distance(self.poses[1], self.poses[2])

        else:
            pose1 = self.blender.interp(self.poses[0], self.poses[1], self.t_last)
            pose2 = self.blender.interp(self.poses[1], self.poses[2], self.t_last)
            self.poses[0] = self.blender.interp(pose1, pose2, self.t_last)
            self.poses[1] = pose2
            self.poses[2] = pose
            self.length_1 = self.blender.distance(self.poses[0], self.poses[1])
            self.length_2 = self.blender.distance(self.poses[1], self.poses[2])

        self.num_poses = min(self.num_poses + 1, 3)
        self.t_last = 0
        self.mutex.release()

    def length(self):
        if self.num_poses == 1:
            return 0
        elif self.num_poses == 2:
            return self.length_1
        return self.length_1 + self.length_2

    def normalize_s(self, s):
        """
        将 s [0, length] 映射到 [0, 1]

        Parameters:
        - s : float, 路程
        """
        total_length = self.length()
        if s >= total_length-1e-6:
            return 1

        t = s / total_length

        # 需要进一步转化为等距离时间标度
        if self.num_poses == 3 and abs(self.length_1 - self.length_2) >= 1e-2:
            t = (np.sqrt(self.length_1**2 + (self.length_2**2-self.length_1**2) * t) -
                     self.length_1) / (self.length_2 - self.length_1)

        return t

    def get_pose(self, s):
        """
        计算二次贝塞尔曲线上的点, s in [0, length]
        """
        with self.mutex:
            if self.num_poses == 1:
                return self.poses[0]

            t = self.normalize_s(s)
            self.t_last = t

            if t == 1:
                self.poses[0] = self.poses[self.num_poses - 1]
                self.num_poses = 1
                self.t_last = 0
                return self.poses[0]

            if self.num_poses == 2:
                pose = self.blender.interp(self.poses[0], self.poses[1], t)
            elif self.num_poses == 3:
                pose1 = self.blender.interp(self.poses[0], self.poses[1], t)
                pose2 = self.blender.interp(self.poses[1], self.poses[2], t)
                pose = self.blender.interp(pose1, pose2, t)

            return pose


class TimeScaler:
    def __init__(self, v0, s_total, vmax, amax, t0=0, s0=0):
        """
        时间标度, 将时间转化成轨迹参数, 并满足限制, 速度连续

        Parameters:
        - v0 : float, 初始速度
        - s : float, 总路程，也即最终位置为 s0 + s
        - vmax : float, 最大速度
        - amax : float, 最大加速度
        - t0 : float, 初始时间
        - s0 : float, 初始位置
        """
        self.update(v0, s_total, vmax, amax, t0, s0)

    def update(self, v0, s_total, vmax, amax, t0=0, s0=0):
        self.t0 = t0
        self.s0 = s0
        self.v0 = v0
        self.vmax = vmax
        self.amax = amax
        self.s_min = v0**2 / 2 / amax
        self.s_total = s_total

        if s_total < self.s_min:
            # print(f's_total is too small, s_min = {self.s_min}, s_total = {s_total}')
            self.amax = v0**2 / 2 / s_total
            # raise ValueError('s_total is too small, s_min = %f' % self.s_min)

        # 不考虑最大速度限制，能达到的最大速度
        vx_2 = self.amax * self.s_total + self.v0 ** 2 / 2
        vx = math.sqrt(vx_2)

        if vx <= self.vmax:
            self.t1 = (vx - self.v0) / self.amax
            self.t2 = 0
            self.t3 = vx / self.amax
            self.vx = vx # 实际能达到的最大速度
        else:
            self.t1 = (self.vmax - self.v0) / self.amax
            self.t2 = (vx_2 - self.vmax ** 2) / self.amax / self.vmax
            self.t3 = self.vmax / self.amax
            self.vx = self.vmax # 实际能达到的最大速度

        self.s1 = self.s0 + self.v0 * self.t1 + self.amax * self.t1 ** 2 / 2  # 加速阶段结束位置
        self.s2 = self.s1 + self.vmax * self.t2  # 匀速阶段结束位置
        self.t_final = self.t1 + self.t2 + self.t3

    def q(self, t):
        t = t - self.t0
        if t <= self.t1:
            return self.s0 + self.v0 * t + self.amax * t ** 2 / 2
        elif t <= self.t1 + self.t2:
            return self.s1 + self.vx * (t - self.t1)
        elif t <= self.t1 + self.t2 + self.t3:
            dt = t - self.t1 - self.t2
            return self.s2 + self.vx * dt - self.amax * dt ** 2 / 2
        else:
            return self.s_total + self.s0

    def dq(self, t):
        t = t - self.t0
        if t <= self.t1:
            return self.v0 + self.amax * t
        elif t <= self.t1 + self.t2:
            return self.vmax
        elif t <= self.t1 + self.t2 + self.t3:
            dt = t - self.t1 - self.t2
            return self.vx - self.amax * dt
        else:
            return 0


class OnlineMotion:

    def __init__(self, origin_pose, vmax=0.04, amax=0.01):
        # t_update: 上一次 TimeScaler 更新时间. 将首次调用 get_pose 视作运动开始时间, 
        #           运动开始后才能更新 t_update, 否则不更新           
        self.t_update = None  

        self.vmax, self.amax = vmax, amax
        self.segment = OnlineSegment(PoseBlender(0.02))
        self.segment.addPose(origin_pose)
        self.time_scaler = TimeScaler(0, 0, vmax, amax)
        self._should_stop = False

    def stop(self):
        self._should_stop = True

    @property
    def _time_scaler_t(self):
        """
        获取 TimeScaler 的时间标度 t, 运动未开始前始终返回 0

        Returns:
        - t : float, 时间标度 t
        """        
        return time.monotonic() - self.t_update if self.t_update else 0

    def add_pose(self, pose):
        with self._mutex:
            pose = pose.copy()
            self.segment.addPose(pose)
            dq = self.time_scaler.dq(self._time_scaler_t)
            self.time_scaler.update(dq, self.segment.length(), self.vmax, self.amax, t0=0)

            if self.t_update is not None:
                self.t_update = time.monotonic()

    def get_pose(self):
        with self._mutex:
            if self.t_update is None:
                self.t_update = time.monotonic()

            s = self.time_scaler.q(self._time_scaler_t)
            return self.segment.get_pose(s)
